#!/bin/bash

# 脚本名称：shrink_disk.sh

# 功能说明：
# 本脚本用于缩小虚拟磁盘或物理磁盘的最后一个分区的大小和文件系统大小。

# 主要流程包括：
# 1. 安装必要的依赖工具（如 qemu-img、qemu-nbd、parted 等）。
# 2. 根据用户输入判断设备是物理磁盘还是虚拟磁盘镜像。
# 3. 如果是虚拟磁盘镜像，则使用 qemu-nbd 挂载到一个空闲的 nbd 设备上。
# 4. 显示设备信息。
# 5. 缩小最后一个 EXT4 分区的大小，并调整文件系统以匹配新容量。
# 6. 断开 NBD 设备，完成操作。

# 注意事项：
# - 仅支持 EXT4 文件系统的分区缩小。
# - 需要 root 权限运行部分命令，如 qemu-nbd、parted、resize2fs。
# - 在执行分区缩小前，建议备份重要数据，以防万一。

# 安装所需的依赖项
install_dependencies() {
    local required=("qemu-img" "qemu-nbd" "partprobe" "resize2fs" "parted")
    local missing=()

    # 检查每个依赖项是否已安装
    for cmd in "${required[@]}"; do
        if ! command -v "$cmd" >/dev/null 2>&1; then
            missing+=("$cmd")
        fi
    done

    # 如果所有依赖项都已安装，则返回
    if [ ${#missing[@]} -eq 0 ]; then
        return 0
    fi

    # 安装缺少的依赖项
    echo "检测到缺少依赖命令: ${missing[*]}，尝试自动安装..."
    sudo apt-get update -qq
    echo "正在安装 qemu-utils parted e2fsprogs..."
    sudo apt-get install -y -qq qemu-utils parted e2fsprogs >/dev/null
}

# 全局变量
disk=""
nbd_device=""
target_device=""
virtual_disk_attached=false
mount_point="/mnt/temp_mount"

# 清理函数，在脚本退出时执行
cleanup() {
    if $virtual_disk_attached && [ -n "$nbd_device" ]; then
        echo "清理：断开 nbd 设备 $nbd_device"
        sudo qemu-nbd -d "$nbd_device"
    fi
    if mountpoint -q "$mount_point"; then
        sudo umount "$mount_point"
        sudo rmdir "$mount_point"
    fi
}
trap cleanup EXIT SIGINT SIGTERM

# 选择一个空闲的 nbd 设备
select_free_nbd() {
    for dev in /dev/nbd*; do
        if [ -z "$(lsblk -n "$dev" | awk 'NR>1')" ]; then
            echo "$dev"
            return 0
        fi
    done
    return 1
}

# 挂载虚拟磁盘
mount_virtual_disk() {
    sudo modprobe nbd max_part=8
    nbd_device=$(select_free_nbd)
    if [ -z "$nbd_device" ]; then
        echo "未找到可用的 nbd 设备，退出。"
        exit 1
    fi
    echo "将虚拟磁盘镜像挂载到 nbd 设备：$nbd_device"
    sudo qemu-nbd -c "$nbd_device" "$disk"
    sudo udevadm settle
    virtual_disk_attached=true
}

# 获取磁盘输入
get_disk_input() {
    if [ -n "$1" ]; then
        disk="$1"
    else
        echo "没有提供设备或镜像文件，现列出所有磁盘："
        lsblk -o NAME,MAJ:MIN,RM,LABEL,SIZE,FSTYPE,UUID,RO,TYPE,MOUNTPOINT,VENDOR,HOTPLUG -e7
        read -rp "请输入设备路径（如 /dev/sda）或虚拟磁盘镜像文件路径: " disk
    fi
}

# 处理磁盘输入
process_disk() {
    if [ -b "$disk" ]; then
        echo "选择的物理磁盘: $disk"
        target_device="$disk"
    elif [ -f "$disk" ]; then
        echo "选择的虚拟磁盘镜像文件: $disk"
        mount_virtual_disk
        target_device="$nbd_device"
    else
        echo "错误： '$disk' 既不是块设备也不是普通文件。"
        exit 1
    fi
}

# 显示磁盘信息
display_disk_info() {
    echo "显示设备信息："
    lsblk -o NAME,MAJ:MIN,RM,LABEL,SIZE,FSTYPE,UUID,RO,TYPE,MOUNTPOINT,VENDOR,HOTPLUG "$target_device"
}

# 缩小最后一个 EXT4 分区
shrink_last_ext4_partition() {
    sudo udevadm settle
    local last_partition=$(lsblk -ln -o NAME,FSTYPE -p "$target_device" | awk '$2=="ext4" {last=$1} END{print last}')
    if [ -z "$last_partition" ]; then
        echo "未找到 EXT4 分区，跳过调整。"
        return
    fi

    echo "找到最后一个 EXT4 分区: $last_partition"
    sudo mkdir -p "$mount_point"
    sudo mount "$last_partition" "$mount_point"
    local used_size=$(sudo df --block-size=1M "$mount_point" | awk 'NR==2 {print $3}')
    sudo umount "$mount_point"
    sudo rmdir "$mount_point"

    local safety_margin=500
    local new_size=$((used_size + safety_margin))
    echo "当前使用空间: ${used_size}M，目标大小: ${new_size}M"

    read -rp "是否继续调整分区大小？(y/n): " confirm
    if [[ "$confirm" != "y" ]]; then
        echo "用户取消了分区调整。"
        return
    fi

    echo "检查文件系统完整性..."
    sudo e2fsck -f -y "$last_partition"

    echo "调整文件系统大小..."
    sudo resize2fs "$last_partition" ${new_size}M

    echo "调整分区大小..."
    local partition_number=$(lsblk -ln -o NAME -p "$target_device" | grep "$last_partition" | sed 's/[^0-9]//g')
    sudo parted "$target_device" resizepart $partition_number ${new_size}MB
}

# 主函数
main() {
    install_dependencies
    get_disk_input "$1"
    process_disk
    display_disk_info
    shrink_last_ext4_partition
}

main "$@"

